import numpy as np
import matplotlib.pyplot as plt
import h5py

# 系统参数配置
class Config:
    RQ = 76.7   # 电阻-品质因数乘积
    Ql = 8893 # 外部品质因数
    beta = 1.2
    Fs = 117.402e6 / 14  # 采样率 [Hz]
    Ts = 1 / Fs  # 采样时间间隔 [s]

# 数据源配置
DATA_FILE = r'E:\MLdata\SimData\coupling_simulation_data.h5'

def read_hdf5_dataset(file_path, set_index):
    """
    从HDF5文件读取指定数据集的数据
    
    Args:
        file_path (str): HDF5文件路径
        set_index (int): 数据集索引
    
    Returns:
        dict: 包含所有电压数据的字典
    """
    with h5py.File(file_path, 'r') as f:
        set_name = f'set_{set_index}'
        if set_name not in f:
            raise ValueError(f"数据集 {set_name} 不存在")
            
        dataset = f[set_name]
        voltages = {
            "Vcav_I": dataset['Vc_real'][:],
            "Vcav_Q": dataset['Vc_imag'][:],
            "Vfor_I": dataset['Vf_star_real'][:],
            "Vfor_Q": dataset['Vf_star_imag'][:],
            "Vrefl_I": dataset['Vr_star_real'][:],
            "Vrefl_Q": dataset['Vr_star_imag'][:],
            "Vf_real": dataset['Vf_real'][:],
            "Vf_imag": dataset['Vf_imag'][:],
            "Vr_real": dataset['Vr_real'][:],
            "Vr_imag": dataset['Vr_imag'][:]
        }
        
        time = f['time'][:]
        
    return voltages, time

class VoltageProcessor:
    @staticmethod
    def clean_voltages(voltages_IQ):
        """清除电压数据中的无效值（NaN或无穷大）"""
        return {key: np.nan_to_num(value) for key, value in voltages_IQ.items()}

    @staticmethod
    def calculate_calibration_matrix(voltages_IQ):
        """计算校准矩阵和校准后的电压值"""
        voltages_IQ = VoltageProcessor.clean_voltages(voltages_IQ)
        
        # 提取IQ分量
        Vcav_I, Vcav_Q = voltages_IQ["Vcav_I"], voltages_IQ["Vcav_Q"]
        Vfor_I, Vfor_Q = voltages_IQ["Vfor_I"], voltages_IQ["Vfor_Q"]
        Vrefl_I, Vrefl_Q = voltages_IQ["Vrefl_I"], voltages_IQ["Vrefl_Q"]
        
        # 构建校准矩阵
        A_matrix = np.column_stack((
            np.append(Vfor_I, Vfor_Q),
            np.append(-Vfor_Q, Vfor_I),
            np.append(Vrefl_I, Vrefl_Q),
            np.append(-Vrefl_Q, Vrefl_I)
        ))
        y_array = np.append(Vcav_I, Vcav_Q)

        try:
            C_array = np.matmul(np.linalg.pinv(A_matrix, rcond=1e-5), y_array.transpose())
            return C_array
        except np.linalg.LinAlgError as e:
            print(f"SVD计算失败: {e}")
            return None

class Plotter:
    @staticmethod
    def plot_voltage_comparison(t_us, Vc_star, Vf_star, Vr_star, X, Y, Vf, Vr):
        """绘制电压比较图表"""
        fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(10, 12))
        
        # 第一个子图：Vc和X*Vf + Y*Vr的比较
        V_combined = X * Vf_star + Y * Vr_star
        V_sum = Vf + Vr
        ax1.plot(t_us, np.abs(Vc_star), label='|Vc|', linewidth=1.5)
        ax1.plot(t_us, np.abs(V_combined), 
                label='|X*Vf_star + Y*Vr_star|', linestyle='--', linewidth=1.5)
        ax1.plot(t_us, np.abs(V_sum), 
                label='|Vf + Vr|', linestyle=':', linewidth=1.5)
        
        # 设置第一个子图的标签和格式
        ax1.set_xlabel(r'Time [$\mu s$]', fontname='Arial', fontsize=12)
        ax1.set_ylabel('Voltage Magnitude', fontname='Arial', fontsize=12)
        ax1.tick_params(labelsize=12)
        ax1.grid(True, linestyle='--', linewidth=0.7)
        ax1.legend(loc='best', fontsize=12, frameon=True, 
                  fancybox=True, edgecolor='black', framealpha=1)
        for spine in ax1.spines.values():
            spine.set_linewidth(1.5)

        # 第二个子图：Vf_star、X*Vf_star和Vf的比较
        X_Vf_star = X * Vf_star
        ax2.plot(t_us, np.abs(Vf_star), label='|Vf_star|', linewidth=1.5)
        ax2.plot(t_us, np.abs(X_Vf_star), 
                label='|X*Vf_star|', linestyle='--', linewidth=1.5)
        ax2.plot(t_us, np.abs(Vf), 
                label='|Vf|', linestyle=':', linewidth=1.5)
        
        # 设置第二个子图的标签和格式
        ax2.set_xlabel(r'Time [$\mu s$]', fontname='Arial', fontsize=12)
        ax2.set_ylabel('Voltage Magnitude', fontname='Arial', fontsize=12)
        ax2.tick_params(labelsize=12)
        ax2.grid(True, linestyle='--', linewidth=0.7)
        ax2.legend(loc='best', fontsize=12, frameon=True, 
                  fancybox=True, edgecolor='black', framealpha=1)
        for spine in ax2.spines.values():
            spine.set_linewidth(1.5)

        # 第三个子图：Vr_star、Y*Vr_star和Vr的比较
        Y_Vr_star = Y * Vr_star
        ax3.plot(t_us, np.abs(Vr_star), label='|Vr_star|', linewidth=1.5)
        ax3.plot(t_us, np.abs(Y_Vr_star), 
                label='|Y*Vr_star|', linestyle='--', linewidth=1.5)
        ax3.plot(t_us, np.abs(Vr), 
                label='|Vr|', linestyle=':', linewidth=1.5)
        
        # 设置第三个子图的标签和格式
        ax3.set_xlabel(r'Time [$\mu s$]', fontname='Arial', fontsize=12)
        ax3.set_ylabel('Voltage Magnitude', fontname='Arial', fontsize=12)
        ax3.tick_params(labelsize=12)
        ax3.grid(True, linestyle='--', linewidth=0.7)
        ax3.legend(loc='best', fontsize=12, frameon=True, 
                  fancybox=True, edgecolor='black', framealpha=1)
        for spine in ax3.spines.values():
            spine.set_linewidth(1.5)
            
        plt.tight_layout()
        plt.show()

def process_file(set_index=0):
    """处理指定索引的数据集"""
    # 读取并处理数据
    try:
        voltages_IQ, time = read_hdf5_dataset(DATA_FILE, set_index)
    except Exception as e:
        print(f"处理数据集 {set_index} 时出错: {e}")
        return
    
    # 计算校准矩阵
    C_array = VoltageProcessor.calculate_calibration_matrix(voltages_IQ)
    if C_array is None:
        return

    # 准备时间数组
    t_us = time * 1e6  # 转换为微秒

    # 提取复数形式的电压值
    Vf_star = voltages_IQ['Vfor_I'] + 1j * voltages_IQ['Vfor_Q']
    Vr_star = voltages_IQ['Vrefl_I'] + 1j * voltages_IQ['Vrefl_Q']
    Vc_star = voltages_IQ['Vcav_I'] + 1j * voltages_IQ['Vcav_Q']
    Vf = voltages_IQ['Vf_real'] + 1j * voltages_IQ['Vf_imag']
    Vr = voltages_IQ['Vr_real'] + 1j * voltages_IQ['Vr_imag']

    # 计算X和Y
    X = complex(C_array[0], C_array[1])
    Y = complex(C_array[2], C_array[3])
    
    print(f'X={X:.2f}, Y={Y:.2f}')

    # 绘制电压幅值比较图
    Plotter.plot_voltage_comparison(t_us, Vc_star, Vf_star, Vr_star, X, Y, Vf, Vr)

def main():
    """主函数"""
    # 处理第一个数据集
    process_file(0)

if __name__ == '__main__':
    main() 